#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
双阶段 YOLO 检测节点

功能说明：
1. 订阅 /yolo_detect 话题接收控制命令（yolostart, yolostop等）
3. 使用双阶段YOLO管道进行检测：
   - 第一阶段：在裁剪图像上检测锚点(数字、锥桶、仪表盘)
   - 第二阶段：对数字ROI进行二次YOLO检测
   - HSV颜色投票：对锥桶ROI进行颜色识别
4. 发布最终结果到/yolo_result话题
5. 发布状态信息到/yolo_status话题

话题接口：
- 订阅：/yolo_detect (std_msgs/String) - 控制命令
- 发布：/yolo_result (std_msgs/String) - 检测结果
- 发布：/yolo_status (std_msgs/String) - 状态信息

检测结果格式：color-number-state (如: red-three-up)
"""
import os
import sys
import gc
import time
import threading
from typing import Optional, Dict, Any

import rospy
import cv2
import numpy as np
from std_msgs.msg import String
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

# 添加yolo-jetson路径到sys.path以导入管道
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, "../../../../"))
yolo_jetson_path = os.path.join(project_root, "yolo-jetson", "detect")
if os.path.exists(yolo_jetson_path):
    sys.path.insert(0, yolo_jetson_path)

try:
    from double_yolo_pipeline import DoubleYoloPipeline
    rospy.loginfo("成功导入 DoubleYoloPipeline")
except ImportError as e:
    rospy.logerr(f"无法导入 DoubleYoloPipeline: {e}")
    rospy.logerr(f"检查路径: {yolo_jetson_path}")
    sys.exit(1)


class YoloDoubleStageNode:
    """双阶段YOLO检测ROS节点"""
    
    def __init__(self):
        rospy.init_node('yolo_double_stage_node', anonymous=True)
        
        # 参数配置
        self.model_path = rospy.get_param('~model_path', 
                                         '/home/jetson/robot_motion/yolo-jetson/weights/all_yolo_all_epochs_0811/best.pt')
        self.device = rospy.get_param('~device', None)  # None表示自动选择
        self.debug = rospy.get_param('~debug', True)
        self.save_dir = rospy.get_param('~save_dir', '/tmp/yolo_double_stage_debug')
        
        # 检测参数
        self.top_cut_ratio = rospy.get_param('~top_cut_ratio', 0.25)
        self.width_center_ratio = rospy.get_param('~width_center_ratio', 0.5)
        self.first_stage_conf = rospy.get_param('~first_stage_conf', 0.6)
        self.second_stage_conf = rospy.get_param('~second_stage_conf', 0.5)
        self.final_conf = rospy.get_param('~final_conf', 0.6)
        self.process_interval = rospy.get_param('~process_interval', 0.1)  # 处理间隔
        
        # 状态管理
        self.is_running = False
        self.bridge = CvBridge()
        self.pipeline: Optional[DoubleYoloPipeline] = None
        self.latest_image: Optional[np.ndarray] = None
        self.frame_count = 0
        self.last_process_time = 0
        self.processing_lock = threading.Lock()
        
        # 统计信息
        self.total_frames_processed = 0
        self.last_result = "None-None-None"
        self.last_processing_time = 0.0
        
        # ROS话题
        self.result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self.status_pub = rospy.Publisher('/yolo_status', String, queue_size=10)
        
        self.command_sub = rospy.Subscriber('/yolo_detect', String, self.command_callback, queue_size=1)
        self.image_sub = rospy.Subscriber('/camera/image_raw', Image, self.image_callback, queue_size=1)
        
        rospy.loginfo("=== 双阶段YOLO检测节点初始化完成 ===")
        rospy.loginfo(f"模型路径: {self.model_path}")
        rospy.loginfo(f"设备: {self.device if self.device else '自动选择'}")
        rospy.loginfo(f"调试模式: {self.debug}")
        rospy.loginfo(f"检测参数: first_conf={self.first_stage_conf}, second_conf={self.second_stage_conf}, final_conf={self.final_conf}")
        rospy.loginfo("等待控制命令...")
        
        # 发布初始状态
        self.publish_status("节点已初始化，等待启动命令")

    def publish_status(self, message: str):
        """发布状态信息"""
        try:
            status_msg = String()
            status_msg.data = f"[DoubleStage] {message}"
            self.status_pub.publish(status_msg)
        except Exception as e:
            rospy.logwarn(f"发布状态失败: {e}")

    def command_callback(self, msg: String):
        """处理控制命令"""
        command = msg.data.strip().lower()
        rospy.loginfo(f"收到命令: {command}")
        
        if command == "yolostart":
            self.start_detection()
        elif command == "yolostop":
            self.stop_detection()
        elif command == "yolostatus":
            self.report_status()
        elif command == "yoloreset":
            self.reset_system()
        else:
            rospy.logwarn(f"未知命令: {command}")
            self.publish_status(f"未知命令: {command}")

    def start_detection(self):
        """启动检测"""
        if self.is_running:
            rospy.loginfo("检测已在运行中")
            self.publish_status("检测已在运行中")
            return
            
        try:
            rospy.loginfo("正在初始化双阶段YOLO管道...")
            self.publish_status("正在初始化YOLO管道...")
            
            if not os.path.exists(self.model_path):
                error_msg = f"模型文件不存在: {self.model_path}"
                rospy.logerr(error_msg)
                self.publish_status(f"错误: {error_msg}")
                return
            
            # 初始化管道
            self.pipeline = DoubleYoloPipeline(
                model_path=self.model_path,
                device=self.device,
                debug=self.debug,
                save_dir=self.save_dir
            )
            
            self.is_running = True
            self.frame_count = 0
            self.total_frames_processed = 0
            self.last_process_time = time.time()
            
            rospy.loginfo("✅ 双阶段YOLO检测已启动")
            self.publish_status("检测已启动，等待图像数据...")
            
        except Exception as e:
            error_msg = f"启动检测失败: {e}"
            rospy.logerr(error_msg)
            self.publish_status(f"启动失败: {e}")
            self.is_running = False
            self.pipeline = None

    def stop_detection(self):
        """停止检测"""
        self.is_running = False
        if self.pipeline:
            # 清理GPU内存
            try:
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except:
                pass
            self.pipeline = None
            gc.collect()
        
        rospy.loginfo("🛑 双阶段YOLO检测已停止")
        self.publish_status(f"检测已停止，共处理 {self.total_frames_processed} 帧")

    def reset_system(self):
        """重置系统"""
        self.stop_detection()
        time.sleep(1)
        self.frame_count = 0
        self.total_frames_processed = 0
        self.last_result = "None-None-None"
        rospy.loginfo("🔄 系统已重置")
        self.publish_status("系统已重置")

    def report_status(self):
        """报告状态"""
        status = "运行中" if self.is_running else "已停止"
        uptime = time.time() - rospy.Time.now().to_sec() if rospy.Time.now().to_sec() > 0 else 0
        
        report = (f"状态: {status} | "
                 f"已处理: {self.total_frames_processed} 帧 | "
                 f"最后结果: {self.last_result} | "
                 f"处理耗时: {self.last_processing_time:.3f}s")
        
        rospy.loginfo(report)
        self.publish_status(report)

    def image_callback(self, msg: Image):
        """处理图像回调"""
        if not self.is_running:
            return
            
        try:
            # 转换图像
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            self.latest_image = cv_image
            self.frame_count += 1
            
            # 检查处理间隔
            current_time = time.time()
            if current_time - self.last_process_time < self.process_interval:
                return
                
            # 非阻塞处理
            if self.processing_lock.acquire(blocking=False):
                try:
                    self.process_current_image()
                finally:
                    self.processing_lock.release()
                    
        except Exception as e:
            rospy.logwarn(f"图像处理回调错误: {e}")

    def process_current_image(self):
        """处理当前图像"""
        if self.latest_image is None or self.pipeline is None:
            return
            
        try:
            start_time = time.time()
            
            # 使用双阶段管道处理图像
            result = self.pipeline.process_image(
                image=self.latest_image.copy(),
                top_cut_ratio=self.top_cut_ratio,
                width_center_ratio=self.width_center_ratio,
                first_stage_conf=self.first_stage_conf,
                second_stage_conf=self.second_stage_conf,
                final_conf=self.final_conf,
                frame_idx=self.total_frames_processed + 1
            )
            
            processing_time = time.time() - start_time
            self.last_processing_time = processing_time
            self.last_process_time = time.time()
            self.total_frames_processed += 1
            
            # 解析结果
            if 'error' in result:
                rospy.logwarn(f"处理错误: {result['error']}")
                self.publish_status(f"处理错误: {result['error']}")
                return
                
            cone_color = result.get('cone_color', 'None')
            number = result.get('number', 'None')
            dashboard_state = result.get('dashboard_state', 'None')
            
            # 构建结果字符串
            final_result = f"{cone_color}-{number}-{dashboard_state}"
            self.last_result = final_result
            
            # 发布结果
            result_msg = String()
            result_msg.data = final_result
            self.result_pub.publish(result_msg)
            
            # 记录日志
            rospy.loginfo(f"[Frame {self.total_frames_processed}] 检测结果: {final_result} "
                         f"(耗时: {processing_time:.3f}s)")
            
            # 定期状态更新
            if self.total_frames_processed % 10 == 0:
                fps = 1.0 / processing_time if processing_time > 0 else 0
                self.publish_status(f"运行中 - 帧率: {fps:.1f} FPS, 最新结果: {final_result}")
                
        except Exception as e:
            rospy.logerr(f"图像处理失败: {e}")
            self.publish_status(f"处理失败: {e}")

    def run(self):
        """运行节点"""
        rospy.loginfo("双阶段YOLO检测节点开始运行...")
        self.publish_status("节点运行中，等待控制命令")
        
        try:
            rospy.spin()
        except KeyboardInterrupt:
            rospy.loginfo("收到停止信号")
        finally:
            self.stop_detection()
            rospy.loginfo("双阶段YOLO检测节点已关闭")


if __name__ == '__main__':
    try:
        node = YoloDoubleStageNode()
        node.run()
    except rospy.ROSInterruptException:
        pass
    except Exception as e:
        rospy.logerr(f"节点启动失败: {e}")
        sys.exit(1)
